#!/usr/bin/python2.5

# Copyright (C) 2007 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Unit tests for the transitfeed_editor module.


import transitfeed_editor
import unittest
import util

from transitfeed import *
from transitfeed_editor import *
from util import RecordingProblemReporter
from StringIO import StringIO

class EditableMemoryZipTestCase(util.TestCaseAsserts):
  """
  TODO:add unit test for Validate function of all GTFSObject
  """
  def setUp(self):
    self.problems = RecordingProblemReporter(self, ("ExpirationDate",))
    self.zipfile = StringIO()
    self.zip = zipfile.ZipFile(self.zipfile, 'a')
    self.zip.writestr(
        "agency.txt",
        "agency_id,agency_name,agency_url,agency_timezone\n"
        "DTA,Demo Agency,http://google.com,America/Los_Angeles\n")
    self.zip.writestr(
        "calendar.txt",
        "service_id,monday,tuesday,wednesday,thursday,friday,saturday,sunday,"
        "start_date,end_date\n"
        "FULLW,1,1,1,1,1,1,1,20070101,20101231\n"
        "WE,0,0,0,0,0,1,1,20070101,20101231\n")
    self.zip.writestr(
        "routes.txt",
        "route_id,agency_id,route_short_name,route_long_name,route_type\n"
        "AB,DTA,,Airport Bullfrog,3\n")
    self.zip.writestr(
        "trips.txt",
        "route_id,service_id,trip_id\n"
        "AB,FULLW,AB1\n")
    self.zip.writestr(
        "stops.txt",
        "stop_id,stop_name,stop_lat,stop_lon\n"
        "BEATTY_AIRPORT,Airport,36.868446,-116.784582\n"
        "BULLFROG,Bullfrog,36.88108,-116.81797\n"
        "STAGECOACH,Stagecoach Hotel,36.915682,-116.751677\n")
    self.zip.writestr(
        "stop_times.txt",
        "trip_id,arrival_time,departure_time,stop_id,stop_sequence\n"
        "AB1,10:00:00,10:00:00,BEATTY_AIRPORT,1\n"
        "AB1,10:20:00,10:20:00,BULLFROG,2\n"
        "AB1,10:25:00,10:25:00,STAGECOACH,3\n")
    self.loader = transitfeed_editor.EditableLoader(
        problems=self.problems,
        extra_validation=True,
        zip=self.zip)

  def appendToZip(self, file, arcname, s):
    """Append s to the arcname in the zip stored in a file object."""
    zip = zipfile.ZipFile(file, 'a')
    zip.writestr(arcname, zip.read(arcname) + s)
    zip.close()

class EditableGenericGTFSObjectTestCase(unittest.TestCase):
  def test__init__(self):
    dicts = {'hello': 100, 'china': '1'}
    gtfs_obj = EditableGenericGTFSObject(dicts)
    self.assertEqual(gtfs_obj.hello, 100)
    self.assertEqual(gtfs_obj.china, '1')
    
class CalendarTestCase(unittest.TestCase):
  def setUp(self):
    attr = {'service_id': 'WCD', 'monday': 4, 'tuesday': 1,
            'wednesday': 0, 'thursday': 3, 'friday': 2,
            'saturday': 0, 'sunday': 1, 'start_date': '19980801', 'end_date': '20081111'}
    self.calendar = Calendar(attr)

  def test_IsValidDate(self):
    self.assertTrue(self.calendar._IsValidDate('19980901'))
    self.assertFalse(self.calendar._IsValidDate('199a0901'))

class CalendarDateTestCase(EditableMemoryZipTestCase):
  def runTest(self):
    """
    TODO: added test for validate
    """

class EditableScheduleTestCase(unittest.TestCase):

  def testInitUsingMemoryDb(self):
    schedule = transitfeed_editor.EditableSchedule(
        None,
        default_problem_reporter,
        True)
    self.assertTrue(schedule._connection)
    self.assertTrue(schedule._cursor)
    self.assertEqual(schedule._original_shapes, {})
    self.assertEqual(schedule._feed_path, None)
    schedule._connection.close()

  def testInitWithFeed(self):
    schedule = transitfeed_editor.EditableSchedule(
                "test/data/good_feed.zip")
    self.assertTrue(schedule._connection)
    self.assertTrue(schedule._cursor)
    self.assertEqual(schedule._original_shapes, {})
    self.assertEqual(schedule._feed_path, "test/data/good_feed.zip")
    schedule._connection.close()

  def testInitWithFileDbWithoutFeed(self):
    schedule = transitfeed_editor.EditableSchedule()
    self.assertTrue(schedule._connection)
    self.assertTrue(schedule._cursor)
    self.assertEqual(schedule._original_shapes, {})
    self.assertEqual(schedule._feed_path, None)
    self.assertTrue('_temp_db_file' in dir(schedule._db)
                 or '_temp_db_filename' in dir(schedule._db))
    schedule._connection.close()

  def testLoad(self):
    schedule = transitfeed_editor.EditableSchedule()
    self.assertFalse(schedule._db.IsGTFSInitialized())
    self.assertFalse('_imported' in dir(schedule))
    schedule.Load("test/data/good_feed.zip")
    self.assertTrue(schedule._db.IsGTFSInitialized())
    self.assertTrue('_imported' in dir(schedule))
    schedule._connection.close()

  def testVerifyExportedShapeId(self):
    feed_filename = "test/data/good_feed.zip"
    schedule = transitfeed_editor.EditableSchedule(feed_filename,
                                                   default_problem_reporter,
                                                   True)
    schedule.Load(feed_filename)
    self.assertTrue(schedule._db.IsTableExist("shapes"))
    
    shapes = schedule._db.GetAll("shapes")
    self.assertEqual(shapes, [])

    trip = schedule.GetTrip("AB2")
    shape_id = trip.shape_id
    shape = schedule.GetShape(shape_id)
    points = ['36.88108,-116.817970000,0.000000000', 
              '36.88357,-116.814500000,0.415085419', 
              '36.88321,-116.814020000,0.473674159', 
              '36.88256,-116.813010000,0.589100287', 
              '36.88072,-116.809460000,0.965760446', 
              '36.87987,-116.807510000,1.163503394', 
              '36.87802,-116.802980000,1.616407738', 
              '36.87767,-116.801970000,1.714421257', 
              '36.87738,-116.801330000,1.779919478', 
              '36.87519,-116.796970000,2.238364032', 
              '36.87458,-116.795910000,2.354644159', 
              '36.87443,-116.795060000,2.432156763', 
              '36.87368,-116.789410000,2.942173498', 
              '36.87336,-116.786670000,3.188760828', 
              '36.87331,-116.786520000,3.203231807', 
              '36.87313,-116.786360000,3.227818696', 
              '36.86935,-116.784360000,3.684748392', 
              '36.86917,-116.784210000,3.708830541', 
              '36.86906,-116.784050000,3.727618269']

    if not shape:
      shape = EditableShape(shape_id)
      schedule.AddGTFSObject(shape, 'shape_id', 
                             EditableShape._SCHEDULE_REFERENCE, None)
    shape.Update(points, schedule)
    schedule.SaveToFeedFile()

    # Generate schedule with the new zip file.
    feed_filename = "test/data/good_feed-new.zip"
    schedule1 = transitfeed_editor.EditableSchedule(feed_filename,
                                                   default_problem_reporter,
                                                   True)
    schedule1.Load(feed_filename)
    self.assertTrue(schedule1._db.IsTableExist("shapes"))
    
    trip1 = schedule1.GetTrip("AB2")
    shape = schedule1.GetShape(trip1.shape_id)
    self.assertEqual(len(shape.points), 19)


class EditableLoaderTestCase(unittest.TestCase):
  def setUp(self):
    schedule = EditableSchedule()
    loader = EditableLoader("test/data/good_feed.zip",
                            schedule,
                            problems=schedule.problem_reporter,
                            extra_validation=False,
                            imported=False,
                            connection=schedule._connection,
                            db=schedule._db)

    self.loader = loader
    self.schedule = schedule
    self._cursor = schedule._cursor

  def test_MatchPattern(self):
    
    self.schedule.AddAgency("DTA1", "http://iflyagency.com", 
                            "America/Los_Angeles", "DTA")
    service_period = self.schedule.GetDefaultServicePeriod()
    
    route = Route(route_id='054C', agency_id='DTA', short_name='10',
                  long_name='Airport - Bullfrog', route_type=3)
    self.schedule.AddRouteObject(route)

    trip = Trip(route=route, service_period=service_period, trip_id='054C-00')
  
    # Import stop_times
    table_name = 'stop_times'
    table_fields = ['trip_id', 'stop_id', 'stop_sequence', 'departure_secs', 
                    'arrival_secs', 'stop_headsign', 'pickup_type', 
                    'drop_off_type', 'shape_dist_traveled']
    table_types = ['char(255)', 'char(255)', 'INT', 'INT', 
                   'INT', 'char(255)', 'char(255)', 'char(255)',
                   'char(255)']
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    
    stop = Stop(36.425288, -117.133162, "Demo Stop 1", "STOP1")
    self.schedule.AddStopObject(stop)
    trip.AddStopTime(stop, schedule=self.schedule, arrival_time="5:11:00", 
                     departure_time="5:12:00")
    stop = Stop(36.423288, -117.143162, "Demo Stop 2", "STOP2")
    self.schedule.AddStopObject(stop)
    trip.AddStopTime(stop, schedule=self.schedule, arrival_time="5:21:00", 
                     departure_time="5:22:00")
    
    self.schedule.AddTripObject(trip)
    
    pattern_id1 = trip.pattern_id
    trip.pattern_id = None
    
    self._cursor.execute("DROP TABLE stop_times;");
    
    # Import stop_times
    table_name = 'stop_times'
    table_fields = ['trip_id', 'stop_id', 'stop_sequence', 'departure_time', 
                    'arrival_time']
    table_types = ['char(255)', 'char(255)', 'char(255)', 'char(255)', 
                   'char(255)']
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    
    rows = [['054C-00', 'STOP1', '1', '5:12:00', '5:11:00']]
    self.loader._db.InsertRows(rows, table_fields, table_name)

    # Import stops
    table_name = 'stops'
    table_fields = ['stop_id', 'stop_lat', 'stop_lon', 'stop_name']
    table_types = ['char(255)', 'char(255)', 'char(255)', 'char(255)']
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    
    rows = [["STOP1", "36.425288", "-117.133262", "Demo Stop 1"],
            ["STOP2", "36.423288", "-117.143162", "Demo Stop 2"]]
    self.loader._db.InsertRows(rows, table_fields, table_name)

    # Import trip
    table_name = 'trips'
    table_fields = ['trip_id', 'route_id', 'shape_id']
    table_types = ['char(255)', 'char(255)', 'char(255)']
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    
    rows = [["054C-00", "054C", ""]]
    self.loader._db.InsertRows(rows, table_fields, table_name)
    
    self.loader._BindTripShapeStopTimesInfo()
    pattern_id2 = trip.pattern_id
    
    # The pattern_id1 is generated with two stop_times, and pattern_id2 is 
    # generated with one stop_time, so they will be different after hashing.
    self.assertNotEqual(pattern_id1, pattern_id2)
    
    table_name = 'stop_times'
    table_fields = ['trip_id', 'stop_id', 'stop_sequence', 'departure_time', 
                    'arrival_time']
    rows = [['054C-00', 'STOP2', '2', '5:22:00', '5:11:00']]
    self.loader._db.InsertRows(rows, table_fields, table_name)

    self.loader._BindTripShapeStopTimesInfo()
    pattern_id2 = trip.pattern_id

    # When the second stop_time is imported into DB, the schedule and the DB
    # will have the same trip, stop_times and stops. And they will generate the
    # same pattern_id.
    self.assertEqual(pattern_id1, pattern_id2)

  def isTableExist(self, table_name):
    sql = """select name from sqlite_master where type = 'table' and
      name = '%s'""" % (table_name)
    h = self._cursor.execute(sql)
    ret = h.fetchall()
    return len(ret) == 1

  def isFieldsExist(self, table_name, fields):
    sql = """select sql from sqlite_master where type = 
      'table' and tbl_name = '%s'""" % (table_name)
    h = self._cursor.execute(sql)
    row = h.fetchone()
    old_fields = re.findall("\((.*)\)", row[0])[0].split(",")
    
    tmp_fields = []
    for row in old_fields:
      tmp_fields.append(row.strip().replace("`", "").split(" ")[0])
    for field in fields:
      if field not in tmp_fields:
        return False
    return True

  def isFieldsTypesMatch(self, table_name, fields, types):
    sql = """select sql from sqlite_master where type = 
      'table' and tbl_name = '%s'""" % (table_name)
    h = self._cursor.execute(sql)
    row = h.fetchone()
    old_fields = re.findall("\((.*)\)", row[0])[0].split(",")
    
    tmp_fields = {}
    for row in old_fields:
      tmp = row.strip().replace("`", "").split(" ")
      tmp_fields[tmp[0].upper()] = tmp[1].upper()
    i = 0
    for field in fields:
      if tmp_fields[field.upper()] != types[i].upper():
        return False
      i += 1
    return True

  def test_CreateDbTable(self):
    table_name = "tmp"
    table_fields = ['abc', 'aaa']
    table_types = ['char(255)', 'integer']
    self.assertFalse(self.isTableExist(table_name))
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    self.assertTrue(self.isTableExist(table_name))
    self.assertTrue(self.isFieldsExist(table_name, table_fields))
    self.assertTrue(self.isFieldsTypesMatch(table_name,
                                            table_fields,
                                            table_types))
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    self.assertTrue(self.isTableExist(table_name))
    self.assertTrue(self.isFieldsExist(table_name, table_fields))
    self.assertTrue(self.isFieldsTypesMatch(table_name,
                                            table_fields,
                                            table_types))

  def test_UpdateDbTable(self):
    table_name = "tmp1"
    table_fields = ['abc', 'aaa']
    table_types = ['char(255)', 'integer']
    table_feed = ["ww", '25']
    self.assertFalse(self.isTableExist(table_name))
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    self.loader._db.UpdateTable(table_feed, table_fields, table_name)
    self.assertTrue(self.isTableExist(table_name))
    sql = "select * from '%s';" % (table_name)
    h = self._cursor.execute(sql)
    row = h.fetchone()
    self.assertTrue(row)
    self.assertEqual(row[0], 'ww')
    self.assertEqual(row[1], 25)

  def test_InsertRows(self):
    table_name = "tmp1"
    table_fields = ['abc', 'aaa']
    table_types = ['char(255)', 'char(10)']
    table_feed = [["test1", 'value1'], ['test2', 'value2']]
    self.assertFalse(self.isTableExist(table_name))
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    self.loader._db.InsertRows(table_feed, table_fields, table_name)
    self.assertTrue(self.isTableExist(table_name))
    sql = "select * from '%s';" % (table_name)
    h = self._cursor.execute(sql)
    row = h.fetchall()
    self.assertTrue(row)
    self.assertEqual(row[0][0], 'test1')
    self.assertEqual(row[0][1], 'value1')
    self.assertEqual(row[1][0], 'test2')
    self.assertEqual(row[1][1], 'value2')
    

  def test_AddExtraColumns(self):
    table_name = "tmp2"
    table_fields = ['abc', 'aaa']
    table_types = ['char(255)', 'integer']
    table_extra_fields = ['a1', 'b1']
    table_extra_types = ['char(256)', 'char(256)']
    self.assertFalse(self.isTableExist(table_name))
    self.loader._db.CreateTable(table_fields, table_types, table_name)
    self.assertTrue(self.isTableExist(table_name))
    self.assertTrue(self.isFieldsExist(table_name, table_fields))
    self.assertFalse(self.isFieldsExist(table_name, table_extra_fields))
    self.loader._db.AddExtraColumns(table_fields, table_extra_fields, 
                                    table_name)
    self.assertTrue(self.isFieldsExist(table_name, table_fields))
    self.assertTrue(self.isFieldsExist(table_name, table_extra_fields))
    self.assertTrue(self.isFieldsTypesMatch(table_name, table_extra_fields,
                                            table_extra_types))

  def test_Load(self):
    self.assertFalse(self.isTableExist(Agency._TABLE_NAME))
    self.assertFalse(self.isTableExist(Route._TABLE_NAME))
    self.assertFalse(self.isTableExist(Trip._TABLE_NAME))
    self.assertFalse(self.isTableExist(Stop._TABLE_NAME))
    self.assertFalse(self.isTableExist(Transfer._TABLE_NAME))
    self.assertFalse(self.isTableExist(Calendar._TABLE_NAME))
    self.assertFalse(self.isTableExist(CalendarDate._TABLE_NAME))
    
    self.loader.Load()
    
    self.assertTrue(self.isTableExist(Agency._TABLE_NAME))
    self.assertTrue(self.isTableExist(Route._TABLE_NAME))
    self.assertTrue(self.isTableExist(Trip._TABLE_NAME))
    self.assertTrue(self.isTableExist(Stop._TABLE_NAME))
    self.assertTrue(self.isTableExist(Transfer._TABLE_NAME))
    self.assertTrue(self.isTableExist(Calendar._TABLE_NAME))
    self.assertTrue(self.isTableExist(CalendarDate._TABLE_NAME))
    
    self.assertTrue(self.isFieldsExist(Agency._TABLE_NAME,
                                       Agency._FIELD_NAMES))
    
    self.assertTrue(self.isFieldsExist(Route._TABLE_NAME,
                                       Route._FIELD_NAMES))
    self.assertTrue(self.isFieldsExist(Trip._TABLE_NAME,
                                       Trip._FIELD_NAMES))
    self.assertTrue(self.isFieldsExist(Stop._TABLE_NAME,
                                       Stop._FIELD_NAMES))
    self.assertTrue(self.isFieldsExist(Transfer._TABLE_NAME,
                                       Transfer._FIELD_NAMES))
    self.assertTrue(self.isFieldsExist(Calendar._TABLE_NAME,
                                       Calendar._FIELD_NAMES))
    self.assertTrue(self.isFieldsExist(CalendarDate._TABLE_NAME,
                                       CalendarDate._FIELD_NAMES))

    self.assertTrue(self.isFieldsTypesMatch(
                      Agency._TABLE_NAME,
                      Agency._FIELD_NAMES,
                      Agency._FIELD_TYPES))
    self.assertTrue(self.isFieldsTypesMatch(
                      Route._TABLE_NAME,
                      Route._FIELD_NAMES,
                      Route._FIELD_TYPES))
    self.assertTrue(self.isFieldsTypesMatch(
                      Trip._TABLE_NAME,
                      Trip._FIELD_NAMES,
                      Trip._FIELD_TYPES))
    self.assertTrue(self.isFieldsTypesMatch(
                      Stop._TABLE_NAME,
                      Stop._FIELD_NAMES,
                      Stop._FIELD_TYPES))
    self.assertTrue(self.isFieldsTypesMatch(
                      Transfer._TABLE_NAME,
                      Transfer._FIELD_NAMES,
                      Transfer._FIELD_TYPES))
    self.assertTrue(self.isFieldsTypesMatch(
                      Calendar._TABLE_NAME,
                      Calendar._FIELD_NAMES,
                      Calendar._FIELD_TYPES))
    self.assertTrue(self.isFieldsTypesMatch(
                      CalendarDate._TABLE_NAME,
                      CalendarDate._FIELD_NAMES,
                      CalendarDate._FIELD_TYPES))
    """test for AddGTFSObject"""
    self.assertTrue('_calendars' in self.schedule.__dict__)

